{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "597a389b",
   "metadata": {},
   "source": [
    "1. load_model 函数\n",
    "功能：加载预训练的图像分割模型。\n",
    "思路图：\n",
    "plaintext\n",
    "开始\n",
    "  ↓\n",
    "实例化 FCNImageSegmenter 模型（类别数为 21）\n",
    "  ↓\n",
    "加载预训练模型参数（从 file_path 指定的文件）\n",
    "  ↓\n",
    "将模型设置为评估模式（关闭 dropout 和 batch normalization）\n",
    "  ↓\n",
    "返回加载好的模型\n",
    "结束\n",
    "分析：\n",
    "该函数负责加载模型权重并初始化模型。\n",
    "在评估模式下，模型不会更新参数，适用于推理任务。\n",
    "\n",
    "\n",
    "2. ImageLoader 类\n",
    "功能：加载图像并进行预处理，使其适合输入到模型中。\n",
    "思路图：\n",
    "plaintext\n",
    "开始\n",
    "  ↓\n",
    "初始化 ImageLoader，定义预处理管道（ToTensor + Normalize）\n",
    "  ↓\n",
    "加载图像文件（从 image_path 指定的路径）\n",
    "  ↓\n",
    "应用预处理管道（转换为张量 + 归一化）\n",
    "  ↓\n",
    "添加批次维度（形状从 (C, H, W) 变为 (1, C, H, W)）\n",
    "  ↓\n",
    "返回预处理后的图像张量\n",
    "结束\n",
    "分析：\n",
    "ToTensor 将图像转换为 PyTorch 张量，并自动将像素值缩放到 [0, 1]。\n",
    "Normalize 使用给定的均值和标准差对图像进行归一化。\n",
    "添加批次维度是为了符合模型的输入要求。\n",
    "\n",
    "\n",
    "3. pred 函数\n",
    "功能：使用模型对输入图像进行预测，得到每个像素的类别索引。\n",
    "思路图：\n",
    "plaintext\n",
    "开始\n",
    "  ↓\n",
    "使用模型进行前向传播（不计算梯度）\n",
    "  ↓\n",
    "获取模型的输出（output[\"out\"]）\n",
    "  ↓\n",
    "对输出取最大值，得到每个像素的预测类别索引\n",
    "  ↓\n",
    "去掉批次维度，返回形状为 (H, W) 的预测结果\n",
    "结束\n",
    "分析：\n",
    "torch.no_grad() 确保在推理过程中不计算梯度，节省内存和计算资源。\n",
    "torch.max(output, dim=1) 在类别维度上取最大值，得到每个像素的预测类别。\n",
    "squeeze(0) 去掉批次维度，使结果形状为 (H, W)。\n",
    "\n",
    "\n",
    "4. find_connected_components 函数\n",
    "功能：在预测结果中找到每个类别的连通区域，并提取每个连通区域的像素坐标。\n",
    "思路图：\n",
    "plaintext\n",
    "开始\n",
    "  ↓\n",
    "遍历每个类别及其对应的名称\n",
    "  ↓\n",
    "创建二值掩码，标记当前类别的像素\n",
    "  ↓\n",
    "使用 connectedComponentsWithStats 找到连通区域\n",
    "  ↓\n",
    "提取每个连通区域的像素坐标\n",
    "  ↓\n",
    "将结果存储到字典中（键为类别名称，值为像素坐标列表）\n",
    "  ↓\n",
    "返回结果字典\n",
    "结束\n",
    "分析：\n",
    "connectedComponentsWithStats 用于找到图像中的连通区域，并返回每个区域的统计信息。\n",
    "通过遍历每个类别，提取其对应的连通区域像素坐标。\n",
    "结果以字典形式返回，便于后续处理和分析。\n",
    "\n",
    "\n",
    "总结\n",
    "load_model：加载模型并设置为评估模式。\n",
    "ImageLoader：加载和预处理图像。\n",
    "pred：使用模型进行预测，得到每个像素的类别索引。\n",
    "find_connected_components：分析预测结果，找到每个类别的连通区域及其像素坐标。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8d64e09",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "from typing import Any, Dict, List\n",
    "\n",
    "from cv2 import connectedComponentsWithStats\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import torch\n",
    "from torch import nn\n",
    "from torchvision import transforms\n",
    "from torchvision.models import resnet\n",
    "from torchvision.models.segmentation.fcn import FCN\n",
    "\n",
    "\n",
    "class FCNHead(nn.Sequential):\n",
    "    def __init__(self, in_channels: int, channels: int) -> None:\n",
    "        inter_channels = in_channels // 4\n",
    "        layers = [\n",
    "            nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(inter_channels),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.1),\n",
    "            nn.Conv2d(inter_channels, channels, 1)\n",
    "        ]\n",
    "\n",
    "        super(FCNHead, self).__init__(*layers)\n",
    "\n",
    "\n",
    "class IntermediateLayerGetter(nn.ModuleDict):\n",
    "\n",
    "    def __init__(self, model: nn.Module, return_layers: Dict[str, str]) -> None:\n",
    "        if not set(return_layers).issubset([name for name, _ in model.named_children()]):\n",
    "            raise ValueError(\"return_layers are not present in model\")\n",
    "\n",
    "        orig_return_layers = return_layers\n",
    "        return_layers = {k: v for k, v in return_layers.items()}\n",
    "        layers = OrderedDict()\n",
    "        for name, module in model.named_children():\n",
    "            layers[name] = module\n",
    "            if name in return_layers:\n",
    "                del return_layers[name]\n",
    "            if not return_layers:\n",
    "                break\n",
    "\n",
    "        super(IntermediateLayerGetter, self).__init__(layers)\n",
    "        self.return_layers = orig_return_layers\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:\n",
    "        out = OrderedDict()\n",
    "        for name, module in self.named_children():\n",
    "            x = module(x)\n",
    "            if name in self.return_layers:\n",
    "                out_name = self.return_layers[name]\n",
    "                out[out_name] = x\n",
    "        return out\n",
    "\n",
    "\n",
    "class FCNImageSegmenter(FCN):\n",
    "    def __init__(self, num_classes: int = 21, **kwargs: Any) -> None:\n",
    "        backbone = resnet.resnet101(\n",
    "            weights=None,\n",
    "            replace_stride_with_dilation=[False, True, True],\n",
    "        )\n",
    "        return_layers = {\"layer4\": \"out\"}\n",
    "        return_layers[\"layer3\"] = \"aux\"\n",
    "        backbone = IntermediateLayerGetter(\n",
    "            backbone, return_layers=return_layers)\n",
    "\n",
    "        inplanes = 1024\n",
    "        aux_classifier = FCNHead(inplanes, num_classes)\n",
    "        inplanes = 2048\n",
    "        classifier = FCNHead(inplanes, num_classes)\n",
    "\n",
    "        super(FCNImageSegmenter, self).__init__(\n",
    "            backbone, classifier, aux_classifier)\n",
    "\n",
    "\n",
    "def load_model(file_path: str) -> FCNImageSegmenter:\n",
    "    # 实例化 FCNImageSegmenter 模型，设置类别数为 21\n",
    "    model = FCNImageSegmenter(num_classes=21)\n",
    "    # 加载预训练模型参数，file_path 是模型权重文件的路径\n",
    "    model.load_state_dict(torch.load(file_path))\n",
    "    # 将模型设置为评估模式，关闭 dropout 和 batch normalization 等训练时使用的特性\n",
    "    model.eval()\n",
    "    return model\n",
    "\n",
    "class ImageLoader:\n",
    "    def __init__(self, mean: List[float] = [0.485, 0.456, 0.406], std: List[float] = [0.229, 0.224, 0.225]) -> None:\n",
    "        # 定义图像预处理管道，包括转换为张量和归一化\n",
    "        self.transform = transforms.Compose([\n",
    "            transforms.ToTensor(),  # 将图像转换为 PyTorch 张量\n",
    "            transforms.Normalize(mean=mean, std=std)  # 使用给定的均值和标准差进行归一化\n",
    "        ])\n",
    "    \n",
    "    def load_image(self, image_path: str) -> torch.Tensor:\n",
    "        # 加载图像文件\n",
    "        image = Image.open(image_path)\n",
    "        # 应用预处理管道，将图像转换为张量并归一化\n",
    "        image_tensor = self.transform(image)\n",
    "        # 添加批次维度，形状从 (C, H, W) 变为 (1, C, H, W)\n",
    "        image_tensor = image_tensor.unsqueeze(0)\n",
    "        return image_tensor\n",
    "\n",
    "\n",
    "def pred(model: nn.Module, input_batch: torch.Tensor) -> torch.Tensor:\n",
    "    # 使用模型进行前向传播，不计算梯度\n",
    "    with torch.no_grad():\n",
    "        output = model(input_batch)[\"out\"]\n",
    "    # 获取每个像素的预测类别索引，dim=1 表示在类别维度上取最大值\n",
    "    _, prediction = torch.max(output, dim=1)\n",
    "    # 去掉批次维度，返回形状为 (H, W) 的张量\n",
    "    return prediction.squeeze(0)\n",
    "\n",
    "\n",
    "def find_connected_components(prediction: np.ndarray, label_mapping: Dict[int, str], connectivity: int = 8) -> Dict[str, List[List[List[int]]]]:\n",
    "    result = {}\n",
    "    # 遍历每个类别及其对应的名称\n",
    "    for label, class_name in label_mapping.items():\n",
    "        # 创建二值掩码，标记当前类别的像素\n",
    "        mask = (prediction == label).astype(np.uint8)\n",
    "        # 使用 connectedComponentsWithStats 找到连通区域\n",
    "        _, labels, stats, centroids = connectedComponentsWithStats(mask, connectivity=connectivity)\n",
    "        # 提取每个连通区域的像素坐标\n",
    "        class_points = []\n",
    "        for i in range(1, len(stats)):  # 跳过背景（索引为 0）\n",
    "            # 获取当前连通区域的像素坐标\n",
    "            y, x = np.where(labels == i)\n",
    "            points = list(zip(y, x))  # 将坐标转换为列表\n",
    "            class_points.append(points)\n",
    "        # 将结果存储到字典中，键为类别名称，值为该类别所有连通区域的像素坐标列表\n",
    "        result[class_name] = class_points\n",
    "    return result\n",
    "\n",
    "def main():\n",
    "    label_mapping = {0: 'Background', \n",
    "                     1: 'Aero plane',\n",
    "                     2: 'Bicycle',\n",
    "                     3: 'Bird',\n",
    "                     4: 'Boat',\n",
    "                     5: 'Bottle',\n",
    "                     6: 'Bus',\n",
    "                     7: 'Car',\n",
    "                     8: 'Cat',\n",
    "                     9: 'Chair',\n",
    "                     10: 'Cow',\n",
    "                     11: 'Dining table',\n",
    "                     12: 'Dog',\n",
    "                     13: 'Horse',\n",
    "                     14: 'Motorbike',\n",
    "                     15: 'Person',\n",
    "                     16: 'Potted plant',\n",
    "                     17: 'Sheep',\n",
    "                     18: 'Sofa',\n",
    "                     19: 'Train',\n",
    "                     20: 'Tv/Monitor'}\n",
    "    model = load_model('model.pth')\n",
    "    loader = ImageLoader()\n",
    "    input_batch = loader.load_image('test.jpg')\n",
    "    output = pred(model, input_batch)\n",
    "    components = find_connected_components(output.detach().numpy(), label_mapping)\n",
    "    \n",
    "    for label, class_points in components.items():\n",
    "        for i, points in enumerate(class_points, 1):\n",
    "            print(label, i, points[:10])\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
